#!/usr/bin/env python
# -*- coding:utf-8 -*-
# author: bruanlin
# datetime: 2024/3/1 20:16
# software: PyCharm
# project: [yolov7-pytorch-master]-[change_xml2person]
""" Code Describe :
 nothing!!
"""
import os
# 程序入口
import xml.etree.ElementTree as ET

from tqdm import tqdm

num_person =0
xml_dir = f"D:\CODE\Codeofpycharm\yolov7-pytorch-master\yolov7-pytorch-master\VOC07+12+test\VOCdevkit\VOC2007\Annotations"
xml_save_dir = f"D:\CODE\Codeofpycharm\yolov7-pytorch-master\yolov7-pytorch-master\VOC07+12+test\VOCdevkit\VOC2007\Annotations_new\\"

def traverse_folder(folder_path):
    file_paths = [];
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            file_path = os.path.join(root, file)
            # print(file_path)
            file_paths.append(file_path)
    return file_paths

def change_one_xml_person(xml_file):
    num_person = 0
    # 加载XML文件
    tree = ET.parse(xml_file)
    root = tree.getroot()

    # 查找所有的object元素
    objects = root.findall("./object")
    for obj in objects:
        # 在每个object元素中找到name元素的文本
        name_element = obj.find("name")
        print(f"current object is ：{name_element.text}")
        if name_element is not None and name_element.text == "person":
            num_person += 1
            continue
        root.remove(obj)
    print(f"该xml文件中找到person对象数为：{num_person}")
    file_name = xml_file.split("\\")[-1]
    print(file_name)
    # 保存修改后的XML文件
    print(xml_save_dir)
    save_path = os.path.join(xml_save_dir, file_name)
    tree.write(save_path)

def change_xml_person():
    file_paths = traverse_folder(xml_dir)
    for file_paths in tqdm(file_paths):
        change_one_xml_person(file_paths)

if __name__ == '__main__':
    # change_one_xml_person(f"D:\CODE\Codeofpycharm\yolov7-pytorch-master\yolov7-pytorch-master\VOC07+12+test\VOCdevkit\VOC2007\Annotations\\000001.xml")
    change_xml_person()